%----------------------------------------------
% Formatting and file path
%----------------------------------------------

clear;
clc;
format bank;

%----------------------------------------------
% Read in data
%----------------------------------------------

load('../../Data/structural/output/Structural_Results_at_ss_assets_minus_land_of_rich.mat')

av_pce=11668.7;

misallocation_summary=zeros(100,1);
misallocation_num_summary=zeros(100,1);

J=100;

for j=1:J

%----------------------------------------------
% Simplify terms
%----------------------------------------------

k_s=k_bl_pre_trans+(av_pce*(0.05*j));
fk_s = a.*k_s.^2 + b.*k_s;


%----------------------------------------------
% Estimate h*, l*, h' at simulated capital level
%----------------------------------------------

psi_l_sim = ratio_psi_l.*psi_l;
psi_h_sim = ratio_psi_l.*psi_h;

l_opt_s = zeros(N,1);
h_opt_s = zeros(N,1);
hiredin_opt_s = zeros(N,1);
profit_opt_s = zeros(N,1);

for i=1:N
    profit = @(x) -A(i)*fk_s(i)*(x(1) + x(3))^beta - w_bl(i)*x(2) + phi*w_bl(i)*x(3) + 0.5*(sqrt(psi_l_sim(i))*x(1) + sqrt(psi_h_sim(i))*x(2))^2;
    [argmax, max] = fmincon(profit, [0,0,0], [1,1,0], R, [], [], [0,0,0], [R,H(i),Hire_in_max]);
    l_opt_s(i) = argmax(1);
    h_opt_s(i) = argmax(2);
    hiredin_opt_s(i) = argmax(3);
    profit_opt_s(i) = -max;
end

clear max argmax

% Check if it would be optimal to liquidate capital:
for i=1:N
    profitalt = @(x) -rho*k_s(i) - w_bl(i)*x + 0.5*psi_h_sim(i)*x^2;
    [argmax, max] = fmincon(profitalt, 0, [], [], [], [], 0, H(i));
    if -max > profit_opt_s(i)
        l_opt_s(i) = 0;
        hiredin_opt_s(i) = 0;
        h_opt_s(i) = argmax;
        profit_opt_s(i) = -max;
    end
end

clear max argmax

% Trim :
l_opt_s(l_opt_s < 1) = 0;
h_opt_s(h_opt_s < 1) = 0;
hiredin_opt_s(hiredin_opt_s < 1) = 0;

% Optimal cases:
case_opt_s = zeros(N,1);

for i=1:N
    if l_opt_s(i) > 0 && h_opt_s(i) > 0 && hiredin_opt_s(i) > 0
        case_opt_s(i) = 1;
    elseif l_opt_s(i) > 0 && h_opt_s(i) > 0 && hiredin_opt_s(i) == 0
        case_opt_s(i) = 2;
    elseif l_opt_s(i) > 0 && h_opt_s(i) == 0 && hiredin_opt_s(i) > 0
        case_opt_s(i) = 3;
    elseif l_opt_s(i) > 0 && h_opt_s(i) == 0 && hiredin_opt_s(i) == 0
        case_opt_s(i) = 4;
    elseif l_opt_s(i) == 0 && h_opt_s(i) > 0 && hiredin_opt_s(i) > 0
        case_opt_s(i) = 5;
    elseif l_opt_s(i) == 0 && h_opt_s(i) > 0 && hiredin_opt_s(i) == 0
        case_opt_s(i) = 6;
    elseif l_opt_s(i) == 0 && h_opt_s(i) == 0
        case_opt_s(i) = 7;
    end
end


%----------------------------------------------
% Estimate value of misallocation at simulated capital level
%----------------------------------------------

misallocation_s = zeros(N,1);

for i=1:N
    if case_opt(i) == 5 || case_opt(i) == 6   % Misallocation is zero if case 5 or 6 dominates at high SS
        misallocation_s(i) = 0;
    else
        misallocation_s(i) = profit_opt(i) - profit_opt_s(i);
    end
end

misallocation_s(misallocation_s < 0) = 0;

misallocation_untopcoded_s=misallocation_s;
total_misallocation_untopcoded_s=sum(misallocation_untopcoded_s);
num_misallocated_untopcoded_s = sum(misallocation_untopcoded_s(:) ~= 0);

% Top-code top 5% outliers
misallocation_95_pcile_s = prctile(misallocation_s,95);
misallocation_s(misallocation_s>misallocation_95_pcile_s) = misallocation_95_pcile_s;
total_misallocation_s=sum(misallocation_s);

% Calculate total number of individuals misallocated
num_misallocated_s = sum(misallocation_s(:) ~= 0);

misallocation_summary(j)=total_misallocation_s;
misallocation_num_summary(j)=num_misallocated_s;

% Stop loop if value of misallocation = 0

    if total_misallocation_s==0
        break
    end

end

%----------------------------------------------
% Plot changing misallocation and number misallocated as capital increases
%----------------------------------------------

misallocation_summary(j+1:J) = [];
misallocation_num_summary(j+1:J) = [];
pc_of_pce=5:5:j*5;

total_transfer_amt = N .* av_pce .* (pc_of_pce./100);

yyaxis left
plot(pc_of_pce,total_transfer_amt);
xlabel('Value of transfer (% of STUP PCE)')
ylabel('Total value of misallocation (BDT)')
plot(pc_of_pce,misallocation_summary);
hold on
yyaxis right
plot(pc_of_pce,misallocation_num_summary);
ylabel('Number of individuals misallocated')

hold on

yyaxis left
plot(pc_of_pce,misallocation_summary);
xlabel('Value of transfer (% of STUP PCE)')
ylabel('Total value of misallocation (solid) and transfers (dashed) (BDT)')
plot(pc_of_pce,total_transfer_amt);

yyaxis right
plot(pc_of_pce,misallocation_num_summary);
ylabel('Number of individuals misallocated')

hold off

% Plot without  total number misallocated

yyaxis left
plot(pc_of_pce,total_transfer_amt);
xlabel('Value of transfer (% of STUP PCE)')
ylabel('Total value of transfers (BDT)')
set(gca,'FontSize',20)

yyaxis right
plot(pc_of_pce,misallocation_summary);
ylabel('Total value of misallocation (BDT)')
set(gca,'FontSize',20)

save('../../Data/structural/intermediate/counterfactuals/Simulate_increase_k_by_multiples_of_avg_pce_results')

T=table(pc_of_pce',misallocation_summary,misallocation_num_summary,total_transfer_amt');
filename='../../Data/structural/intermediate/counterfactuals/Simulate_increase_k_by_multiples_of_avg_pce_results.csv';

writetable(T,filename);
